#!/usr/bin/env python
from typing import List, Tuple

from langchain_core.documents import Document

from service.reranker import reranker


class HybridSearch:
    """
    基于RRF算法，合并多种检索系统的检索结果，当前支持bm25(es)检索及向量库检索结果的合并，要求，传入的rank满足格式: [('xxx', 11.2),... ('xxx', 56.3)]
    """
    def __init__(self):
        self.reranker = reranker

    def get_difference(self, match_result: List[Document], vector_result: List[Document]) \
            -> List[Document]:
        # 获取vector_result - match_result差集
        match_contents = {doc.page_content for doc in match_result}
        # 直接在 vector_result 上筛选出不在 match_contents 中的文档
        difference_docs = [doc for doc in vector_result if doc.page_content not in match_contents]
        return difference_docs

    def get_intersection(self, match_result: List[Document], vector_result: List[Document]) \
            -> List[Document]:
        # 获取vector_result ∩ match_result交集，
        match_contents = {doc.page_content for doc in match_result}
        # 直接在 vector_result 上筛选出在 match_contents 中的文档
        intersection_docs = [doc for doc in vector_result if doc.page_content in match_contents]
        return intersection_docs

    def model_rerank(self, query: str, match_result: List[Document], vector_result: List[Document],
                     final_top_k: int = 10, reranker_threshold: float = 0.5
                     ) -> List[Tuple[Document, float]]:

        same_contents = self.get_intersection(match_result[:10], vector_result)
        # 获取两路召回中的差集
        deputed_data = self.get_difference(match_result, vector_result)
        # 重排序
        common_reranked_res = self.reranker.rerank(query, deputed_data, top_k=9)
        same_reranked_res = self.reranker.rerank(query, same_contents, top_k=3)
        # 合并结果，优先插入相同内容的结果
        final_results = []  # 其中的元素已经是Tuple了
        for doc in same_reranked_res:
            if doc not in common_reranked_res:
                final_results.append(doc)
        # 添加其余的重排序结果
        for doc in common_reranked_res:
            if doc not in final_results:
                final_results.append(doc)
        # print(final_results)
        # 应用阈值过滤
        return [doc for doc in final_results if doc[1] > reranker_threshold][:final_top_k]


if __name__ == '__main__':
    # re = ReRanker()
    # print('模型加载完成')
    pass
